// Copyright Envoy AI Gateway Authors
// SPDX-License-Identifier: Apache-2.0
// The full text of the Apache license is available in the LICENSE file at
// the root of the repo.

package e2e

import (
	"cmp"
	"context"
	"io"
	"net/http"
	"os"
	"strconv"
	"strings"
	"testing"
	"time"

	"github.com/stretchr/testify/require"

	"github.com/envoyproxy/ai-gateway/tests/internal/e2elib"
)

// Test_Examples_ProviderFallback tests the basic example in the examples/provider_fallback directory.
func Test_Examples_ProviderFallback(t *testing.T) {
	const baseManifest = "../../examples/provider_fallback/base.yaml"
	const fallbackManifest = "../../examples/provider_fallback/fallback.yaml"
	const egSelector = "gateway.envoyproxy.io/owning-gateway-name=provider-fallback"

	// This requires the following environment variables to be set:
	//   - TEST_AWS_ACCESS_KEY_ID
	//   - TEST_AWS_SECRET_ACCESS_KEY
	//
	// A test case will be skipped if the corresponding environment variable is not set.
	read, err := os.ReadFile(baseManifest)
	require.NoError(t, err)
	// Replace the placeholder with the actual credentials.
	awsAccessKeyID := os.Getenv("TEST_AWS_ACCESS_KEY_ID")
	awsSecretAccessKey := os.Getenv("TEST_AWS_SECRET_ACCESS_KEY")
	if awsAccessKeyID == "" || awsSecretAccessKey == "" {
		t.Skip("skipped due to missing credentials")
	}

	replaced := strings.ReplaceAll(string(read), "AWS_ACCESS_KEY_ID", cmp.Or(awsAccessKeyID, "dummy-aws-access-key-id"))
	replaced = strings.ReplaceAll(replaced, "AWS_SECRET_ACCESS_KEY", cmp.Or(awsSecretAccessKey, "dummy-aws-secret-access-key"))
	require.NoError(t, e2elib.KubectlApplyManifestStdin(t.Context(), replaced))
	e2elib.RequireWaitForGatewayPodReady(t, egSelector)
	t.Cleanup(func() {
		_ = e2elib.KubectlDeleteManifest(context.Background(), baseManifest)
		_ = e2elib.KubectlDeleteManifest(context.Background(), fallbackManifest)
	})

	const body = `{"model": "us.meta.llama3-2-1b-instruct-v1:0","messages": [{"role": "user", "content": "Say this is a test!"}],"temperature": 0.7}`

	// At this point, we haven't configured the fallback for the HTTPRoute generated by the AI Gateway.
	// So, no matter how many times we try, we should always get a 503 error.
	for i := range 5 {
		t.Run("no-fallback/"+strconv.Itoa(i), func(t *testing.T) {
			fwd := e2elib.RequireNewHTTPPortForwarder(t, e2elib.EnvoyGatewayNamespace, egSelector, e2elib.EnvoyGatewayDefaultServicePort)
			defer fwd.Kill()

			ctx, cancel := context.WithTimeout(t.Context(), 10*time.Second)
			defer cancel()

			req, err := http.NewRequestWithContext(ctx, http.MethodPost, fwd.Address()+"/v1/chat/completions", strings.NewReader(body))
			require.NoError(t, err)
			req.Header.Set("Content-Type", "application/json")
			resp, err := http.DefaultClient.Do(req)
			require.NoError(t, err)
			defer resp.Body.Close()
			respBody, err := io.ReadAll(resp.Body)
			require.NoError(t, err)
			t.Logf("Response body: %s", string(respBody))
			require.Equal(t, http.StatusServiceUnavailable, resp.StatusCode)
			require.Contains(t, string(respBody), "upstream connect error or disconnect/reset before headers")
		})
	}

	// Now let's apply the fallback configuration.
	require.NoError(t, e2elib.KubectlApplyManifest(t.Context(), fallbackManifest))
	// Wait for the fallback configuration to be applied.
	time.Sleep(5 * time.Second)

	// At this point, the fallback configuration should be applied. So the request must be either
	// successful or return a 401 error due to the secret key not being propagated yet.
	require.Eventually(t, func() bool {
		fwd := e2elib.RequireNewHTTPPortForwarder(t, e2elib.EnvoyGatewayNamespace, egSelector, e2elib.EnvoyGatewayDefaultServicePort)
		defer fwd.Kill()

		ctx, cancel := context.WithTimeout(t.Context(), 10*time.Second)
		defer cancel()

		req, err := http.NewRequestWithContext(ctx, http.MethodPost, fwd.Address()+"/v1/chat/completions", strings.NewReader(body))
		require.NoError(t, err)
		req.Header.Set("Content-Type", "application/json")
		resp, err := http.DefaultClient.Do(req)
		require.NoError(t, err)
		defer resp.Body.Close()
		respBody, err := io.ReadAll(resp.Body)
		require.NoError(t, err)
		t.Logf("Response body: %s", string(respBody))
		if resp.StatusCode != http.StatusOK && resp.StatusCode != http.StatusUnauthorized {
			// This means that the fallback configuration is not working as expected.
			t.Fatalf("Unexpected status code: %d: %s", resp.StatusCode, string(respBody))
		}
		t.Logf("Response status code: %d", resp.StatusCode)
		return resp.StatusCode == http.StatusOK
	}, time.Second*60, time.Second*5, "Waiting for the fallback configuration to be applied")

	// Now we should be able to get a response from the fallback provider (AWS) without dropping any requests.
	for i := range 5 {
		t.Run("with-fallback/"+strconv.Itoa(i), func(t *testing.T) {
			fwd := e2elib.RequireNewHTTPPortForwarder(t, e2elib.EnvoyGatewayNamespace, egSelector, e2elib.EnvoyGatewayDefaultServicePort)
			defer fwd.Kill()

			ctx, cancel := context.WithTimeout(t.Context(), 10*time.Second)
			defer cancel()

			req, err := http.NewRequestWithContext(ctx, http.MethodPost, fwd.Address()+"/v1/chat/completions", strings.NewReader(body))
			require.NoError(t, err)
			req.Header.Set("Content-Type", "application/json")
			resp, err := http.DefaultClient.Do(req)
			require.NoError(t, err)
			defer resp.Body.Close()
			respBody, err := io.ReadAll(resp.Body)
			require.NoError(t, err)
			t.Logf("Response body: %s", string(respBody))
			require.Equal(t, http.StatusOK, resp.StatusCode)
		})
	}
}
